package drds.common.db.xml_mapping.executor;


import drds.common.db.xml_mapping.executor.key_generator.AutoSetModePrimaryKeyGenerator;
import drds.common.db.xml_mapping.executor.key_generator.NoPrimaryKeyGenerator;
import drds.common.db.xml_mapping.executor.key_generator.PrimaryKeyGenerator;
import drds.common.db.xml_mapping.executor.statement.StatementHandler;
import drds.common.db.xml_mapping.mapping.BindSql;
import drds.common.db.xml_mapping.mapping.Mapper;
import drds.common.db.xml_mapping.session.Configuration;
import drds.common.db.xml_mapping.session.ResultHandler;
import drds.common.db.xml_mapping.session.RowBounds;
import drds.common.db.xml_mapping.transaction.Transaction;

import java.sql.BatchUpdateException;
import java.sql.Connection;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;


public class BatchExecutor extends BaseExecutor {

  public static final int BATCH_UPDATE_RETURN_VALUE = Integer.MIN_VALUE + 1002;

  private final List<Statement> statementList = new ArrayList<>();
  private final List<BatchResult> batchResultList = new ArrayList<>();
  private String currentSql;
  private Mapper currentMapper;

  public BatchExecutor(Configuration configuration, Transaction transaction) {
    super(configuration, transaction);
  }

  @Override
  public int doUpdate(Mapper mapper, Object parameterObject) throws SQLException {
    final Configuration configuration = mapper.getConfiguration();
    final StatementHandler statementHandler = configuration.newStatementHandler(this, mapper, parameterObject, RowBounds.DEFAULT, null, null);
    final BindSql bindSql = statementHandler.getBindSql();
    final String sql = bindSql.getSql();
    final Statement statement;
    if (sql.equals(currentSql) && mapper.equals(currentMapper)) {
      int last = statementList.size() - 1;
      statement = statementList.get(last);

      statementHandler.parameterize(statement);//fix Issues 322
      BatchResult batchResult = batchResultList.get(last);
      batchResult.addParameterObject(parameterObject);
    } else {
      Connection connection = getConnection();
      statement = statementHandler.prepare(connection, transaction.getTimeout());
      statementHandler.parameterize(statement);    //fix Issues 322
      currentSql = sql;
      currentMapper = mapper;
      statementList.add(statement);
      batchResultList.add(new BatchResult(mapper, sql, parameterObject));
    }
    statementHandler.batch(statement);
    return BATCH_UPDATE_RETURN_VALUE;
  }

  @Override
  public <E> List<E> doQuery(Mapper mapper, Object parameterObject, RowBounds rowBounds, ResultHandler resultHandler, BindSql bindSql)
    throws SQLException {
    Statement stmt = null;
    try {
      flushStatements();
      Configuration configuration = mapper.getConfiguration();
      StatementHandler handler = configuration.newStatementHandler(executor, mapper, parameterObject, rowBounds, resultHandler, bindSql);
      Connection connection = getConnection();
      stmt = handler.prepare(connection, transaction.getTimeout());
      handler.parameterize(stmt);
      return handler.query(stmt, resultHandler);
    } finally {
      closeStatement(stmt);
    }
  }


  @Override
  public List<BatchResult> doFlushStatements(boolean isRollback) throws SQLException {
    try {
      List<BatchResult> results = new ArrayList<>();
      if (isRollback) {
        return Collections.emptyList();
      }
      for (int i = 0, n = statementList.size(); i < n; i++) {
        Statement statement = statementList.get(i);

        BatchResult batchResult = batchResultList.get(i);
        try {
          batchResult.setUpdateCounts(statement.executeBatch());
          Mapper MappedStatement = batchResult.getMappedStatement();
          List<Object> parameterObjects = batchResult.getParameterObjects();
          PrimaryKeyGenerator primaryKeyGenerator = MappedStatement.getPrimaryKeyGenerator();
          if (AutoSetModePrimaryKeyGenerator.class.equals(primaryKeyGenerator.getClass())) {
            AutoSetModePrimaryKeyGenerator autoSetModePrimaryKeyGenerator = (AutoSetModePrimaryKeyGenerator) primaryKeyGenerator;
            autoSetModePrimaryKeyGenerator.processBatch(MappedStatement, statement, parameterObjects);
          } else if (!NoPrimaryKeyGenerator.class.equals(primaryKeyGenerator.getClass())) { //issue #141
            for (Object parameter : parameterObjects) {
              primaryKeyGenerator.generate(this, MappedStatement, statement, parameter);
            }
          }
          // Close statement to close cursor
          closeStatement(statement);
        } catch (BatchUpdateException e) {
          StringBuilder message = new StringBuilder();
          message.append(batchResult.getMappedStatement().getId())
            .append(" (batch index #")
            .append(i + 1)
            .append(")")
            .append(" failed.");
          if (i > 0) {
            message.append(" ")
              .append(i)
              .append(" prior sub executor(s) completed successfully, but will be rolled back.");
          }
          throw new BatchExecutorException(message.toString(), e, results, batchResult);
        }
        results.add(batchResult);
      }
      return results;
    } finally {
      for (Statement stmt : statementList) {
        closeStatement(stmt);
      }
      currentSql = null;
      statementList.clear();
      batchResultList.clear();
    }
  }

}
